module latlon_process_mod

  use mpi
  use flogger
  use const_mod
  use latlon_mesh_mod
  use latlon_halo_mod

  implicit none

  private

  public process_init
  public process_final
  public proc
  public latlon_zonal_circle_type

  type latlon_process_neighbor_type
    integer :: id       = MPI_PROC_NULL
    integer :: orient   = 0
    integer :: lon_ibeg = inf_i4
    integer :: lon_iend = inf_i4
    integer :: lat_ibeg = inf_i4
    integer :: lat_iend = inf_i4
  contains
    procedure :: init => latlon_process_neighbor_init
  end type latlon_process_neighbor_type

  type latlon_zonal_circle_type
    integer :: group = MPI_GROUP_NULL
    integer :: comm  = MPI_COMM_NULL
    integer :: np    = 0
    integer :: id    = MPI_PROC_NULL
    integer :: west_ngb_id = MPI_PROC_NULL
    integer :: east_ngb_id = MPI_PROC_NULL
  contains
    procedure :: init => latlon_zonal_circle_init
    final :: latlon_zonal_circle_final
  end type latlon_zonal_circle_type

  type latlon_process_type
    integer np
    integer :: comm           = MPI_COMM_NULL
    integer :: cart_comm      = MPI_COMM_NULL
    integer :: cart_group     = MPI_GROUP_NULL
    integer :: cart_dims(2)   = 0
    integer :: cart_coords(2) = 0
    integer :: id             = MPI_PROC_NULL   ! MPI process ID
    integer num_lon
    integer num_lat
    integer lon_ibeg
    integer lon_iend
    integer lat_ibeg
    integer lat_iend
    type(latlon_zonal_circle_type) zonal_circle
    type(latlon_process_neighbor_type) ngb(4) ! Neighbor processes
    type(latlon_halo_type) halo(4)
    type(latlon_mesh_type) global_mesh
    type(latlon_mesh_type) local_mesh
  contains
    procedure :: init => latlon_process_init
    procedure :: is_root_proc => latlon_process_is_root_proc
    procedure :: stop => latlon_process_stop
    final :: latlon_process_final
  end type latlon_process_type

  type(latlon_process_type), allocatable :: proc

contains

  subroutine process_init()

    allocate(proc)

  end subroutine process_init

  subroutine process_final()

    if (allocated(proc)) deallocate(proc)

  end subroutine process_final

  subroutine latlon_process_init(this, comm, parallel_domain)

    class(latlon_process_type), intent(inout) :: this
    integer, intent(in), optional :: comm
    logical, intent(in), optional :: parallel_domain

    integer ierr

    if (present(comm)) then
      proc%comm = comm
    else
      call MPI_INIT(ierr)
      proc%comm = MPI_COMM_WORLD
    end if

    if (merge(parallel_domain, .true., present(parallel_domain))) then
      call setup_mpi_2d(this)
    else
      call setup_mpi_0d(this)
    end if
    call decompose_domain(this)
    call this%zonal_circle%init()

  end subroutine latlon_process_init

  subroutine latlon_process_stop(this, code)

    class(latlon_process_type), intent(inout) :: this
    integer, intent(in) :: code

    integer ierr

    call MPI_ABORT(this%comm, code, ierr)

  end subroutine latlon_process_stop

  subroutine latlon_process_final(this)

    type(latlon_process_type), intent(inout) :: this

    integer ierr

    if (this%cart_group /= MPI_GROUP_NULL) call MPI_GROUP_FREE(this%cart_group, ierr)

    call MPI_FINALIZE(ierr)

  end subroutine latlon_process_final

  pure logical function latlon_process_is_root_proc(this) result(res)

    class(latlon_process_type), intent(in) :: this

    res = this%id == 0

  end function latlon_process_is_root_proc

  subroutine setup_mpi_0d(proc)

    type(latlon_process_type), intent(inout) :: proc

    integer ierr

    call MPI_COMM_SIZE(proc%comm, proc%np, ierr)
    call MPI_COMM_RANK(proc%comm, proc%id, ierr)
    proc%cart_dims = [1, 1]
    proc%cart_coords = [0, 0]
    proc%ngb(:)%id = proc%id

  end subroutine setup_mpi_0d

  subroutine setup_mpi_2d(proc)

    type(latlon_process_type), intent(inout) :: proc

    integer ierr

    call MPI_COMM_SIZE(proc%comm, proc%np, ierr)
    call MPI_COMM_RANK(proc%comm, proc%id, ierr)
    call MPI_DIMS_CREATE(proc%np, 2, proc%cart_dims, ierr)
    call MPI_CART_CREATE(proc%comm, 2, proc%cart_dims, [.true., .false.], .true., proc%cart_comm, ierr)
    call MPI_COMM_GROUP(proc%cart_comm, proc%cart_group, ierr)
    call MPI_CART_COORDS(proc%cart_comm, proc%id, 2, proc%cart_coords, ierr)

    ! Set neighborhood of the process.
    call MPI_CART_SHIFT(proc%cart_comm, 0, 1, proc%ngb(west )%id, proc%ngb(east )%id, ierr)
    call MPI_CART_SHIFT(proc%cart_comm, 1, 1, proc%ngb(south)%id, proc%ngb(north)%id, ierr)

  end subroutine setup_mpi_2d

  subroutine decompose_domain(proc)

    type(latlon_process_type), intent(inout) :: proc

    integer ierr, tmp_id(1), i, j
    integer hw

    ! Set initial values for num_lon, num_lat, lon_ibeg, lat_ibeg.
    proc%num_lon  = proc%global_mesh%num_lon
    proc%num_lat  = proc%global_mesh%num_lat
    proc%lon_ibeg = 1
    proc%lat_ibeg = 1

    call round_robin(proc%cart_dims(1), proc%cart_coords(1), proc%num_lon, proc%lon_ibeg, proc%lon_iend)
    call round_robin(proc%cart_dims(2), proc%cart_coords(2), proc%num_lat, proc%lat_ibeg, proc%lat_iend)

    hw = proc%global_mesh%lon_halo_width

    ! TODO: Support 2D decomposition.
    call proc%ngb(west )%init(west , lat_ibeg=proc%lat_ibeg   , lat_iend=proc%lat_iend   )
    call proc%ngb(east )%init(east , lat_ibeg=proc%lat_ibeg   , lat_iend=proc%lat_iend   )
    call proc%ngb(south)%init(south, lon_ibeg=proc%lon_ibeg-hw, lon_iend=proc%lon_iend+hw)
    call proc%ngb(north)%init(north, lon_ibeg=proc%lon_ibeg-hw, lon_iend=proc%lon_iend+hw)

    ! Setup local mesh.
    call proc%local_mesh%init_from_parent(proc%global_mesh               , &
                                          proc%global_mesh%lon_halo_width, &
                                          proc%global_mesh%lat_halo_width, &
                                          proc%lon_ibeg, proc%lon_iend   , &
                                          proc%lat_ibeg, proc%lat_iend)

    ! Setup halos.
    do i = 1, size(proc%ngb)
      select case (proc%ngb(i)%orient)
      case (west, east)
        call proc%halo(i)%init(proc%local_mesh, proc%ngb(i)%orient, MPI_DOUBLE,          &
                               host_id=proc%id, ngb_proc_id=proc%ngb(i)%id,              &
                               lat_ibeg=proc%ngb(i)%lat_ibeg, lat_iend=proc%ngb(i)%lat_iend)
      case (south, north)
        call proc%halo(i)%init(proc%local_mesh, proc%ngb(i)%orient, MPI_DOUBLE,          &
                               host_id=proc%id, ngb_proc_id=proc%ngb(i)%id,              &
                               lon_ibeg=proc%ngb(i)%lon_ibeg, lon_iend=proc%ngb(i)%lon_iend)
      end select
    end do

  end subroutine decompose_domain

  subroutine latlon_process_neighbor_init(this, orient, lon_ibeg, lon_iend, lat_ibeg, lat_iend)

    class(latlon_process_neighbor_type), intent(inout) :: this
    integer, intent(in) :: orient
    integer, intent(in), optional :: lon_ibeg
    integer, intent(in), optional :: lon_iend
    integer, intent(in), optional :: lat_ibeg
    integer, intent(in), optional :: lat_iend

    this%orient = orient

    select case (orient)
    case (west, east)
      this%lat_ibeg = lat_ibeg
      this%lat_iend = lat_iend
    case (south, north)
      this%lon_ibeg = lon_ibeg
      this%lon_iend = lon_iend
    end select

  end subroutine latlon_process_neighbor_init

  subroutine latlon_zonal_circle_init(this)

    class(latlon_zonal_circle_type), intent(inout) :: this

    integer ierr, i, num_lon, ibeg, iend
    integer west_cart_id, east_cart_id, tmp_id(1)
    integer, allocatable :: zonal_proc_id(:)

    if (proc%cart_comm == MPI_COMM_NULL) then
      this%comm = proc%comm
      this%west_ngb_id = this%id
      this%east_ngb_id = this%id
    else
      allocate(zonal_proc_id(proc%cart_dims(1)))
      do i = 1, proc%cart_dims(1)
        call MPI_CART_RANK(proc%cart_comm, [i-1,proc%cart_coords(2)], zonal_proc_id(i), ierr)
      end do
      call MPI_GROUP_INCL(proc%cart_group, size(zonal_proc_id), zonal_proc_id, this%group, ierr)
      call MPI_COMM_CREATE_GROUP(proc%cart_comm, this%group, sum(zonal_proc_id), this%comm, ierr)
      call MPI_COMM_SIZE(this%comm, this%np, ierr)
      call MPI_COMM_RANK(this%comm, this%id, ierr)
      deallocate(zonal_proc_id)

      ! Get IDs of the west and east neighbors in zonal circle comm.
      call MPI_CART_SHIFT(proc%cart_comm, 0, 1, west_cart_id, east_cart_id, ierr)
      call MPI_GROUP_TRANSLATE_RANKS(proc%cart_group, 1, [west_cart_id], this%group, tmp_id, ierr); this%west_ngb_id = tmp_id(1)
      call MPI_GROUP_TRANSLATE_RANKS(proc%cart_group, 1, [east_cart_id], this%group, tmp_id, ierr); this%east_ngb_id = tmp_id(1)
    end if

  end subroutine latlon_zonal_circle_init

  subroutine latlon_zonal_circle_final(this)

    type(latlon_zonal_circle_type), intent(inout) :: this

    integer i, k, ierr

    if (this%group /= MPI_GROUP_NULL) call MPI_GROUP_FREE(this%group, ierr)

  end subroutine latlon_zonal_circle_final

  subroutine round_robin(dim, coord, num, ibeg, iend)

    integer, intent(in) :: dim
    integer, intent(in) :: coord
    integer, intent(inout) :: num
    integer, intent(inout) :: ibeg
    integer, intent(inout) :: iend

    integer res_num, tmp_num, i

    res_num = mod(num, dim)
    ibeg = 1
    do i = 0, coord - 1
      if (res_num /= 0 .and. i < res_num) then
        tmp_num = num / dim + 1
      else
        tmp_num = num / dim
      end if
      ibeg = ibeg + tmp_num
    end do
    if (res_num /= 0 .and. coord < res_num) then
      num = num / dim + 1
    else
      num = num / dim
    end if
    iend = ibeg + num - 1

  end subroutine round_robin

end module latlon_process_mod
